# -*- coding: utf-8 -*-
import numpy as np
import pandas as pd
from datetime import datetime as dt
from pandas.plotting import autocorrelation_plot
import matplotlib.pyplot as plt
from statsmodels.tsa.stattools import adfuller
from matplotlib import cm

plt.rcParams['font.sans-serif']=['SimHei']    #画图时使用中文字体
plt.rcParams['axes.unicode_minus'] = False

def plot_scatter(data, temperature):
    '''
    画散点-温度图
    '''
    cmap=cm.coolwarm
    fig, ax = plt.subplots()
    data_num = len(data)
    x = range(data_num)
    ax.scatter(x, y=data, cmap=cmap, c=temperature, s=15)
    ax.set(xlabel='日期', ylabel='电力负荷(MW)',
           title='日期-负荷序列（MW）')
    ax.grid()
    # 将点的颜色设置为与温度有关
    fig.colorbar(cm.ScalarMappable(cmap=cmap))
    fig.savefig(r'../图片/日期-负荷-温度图.png')
    plt.show()
    
def plot_hist(data):
    '''
    画频率密度图
    '''
    fig, ax = plt.subplots()

    # the histogram of the data
    bins = 50
    ax.hist(data, bins, density=True)
    ax.set_xlabel('电力负荷(MW)')
    ax.set_ylabel('频率密度')
    ax.set_title(r'负荷柱状图')
    ax.grid()
    # Tweak spacing to prevent clipping of ylabel
    fig.savefig(r'../图片/频率密度直方图.png')
    fig.tight_layout()
    plt.show()
    
def plot_autocorrelation(data):
    '''
    画自相关图
    '''
    fig = plt.figure()
    autocorrelation_plot(data)

    plt.xlabel('时间延迟')
    plt.ylabel('自相关系数')
    plt.title('自相关图')
    plt.grid()
    fig.savefig(r'../图片/自相关图.png')
    fig.tight_layout()
    plt.show()

def is_weekend(date):
    '''
    判断当天是否属于周末
    '''
    # 返回 date 的星期数（从 0 开始）
    weekno = date.weekday()
    if weekno < 5:
        return False
    else:
        return True

def prepare_data(data, holiday, temperature, year=1997):
    '''
    将数据按天采样，并附上 is_holiday 和当天温度。
    '''
    # 按天采样
    data_day = data.iloc[:, 3:].sum(axis=1)
    # 时间数据处理成合适格式
    data_date = data.iloc[:, 0].astype('str') + '-' + \
            data.iloc[:, 1].astype('str') + '-' +  \
            data.iloc[:, 2].astype('str') 
    # 将日期字符串转为时间戳
    data_date = data_date.apply(lambda x: dt.strptime(x, '%Y-%m-%d'))
    data_day = pd.concat([data_date, data_day], axis=1)
    data_day.columns = ['date', 'energy_load']
    
    # 添加 是否假期 列
    holiday.drop(index=0, axis=0, inplace=True)
    data_day['is_holiday'] = 0
    # 若数据的日期在假期表格内
    holiday_rows = data_day.loc[data_day['date'].isin(list(holiday.iloc[:,year-1997]))].index
    data_day.iloc[holiday_rows, 2] = 1
    
    
    # 添加 是否周末 列
    data_day['is_weekend'] =0
    weekend_rows = data_day.loc[data_day['date'].apply(lambda x: is_weekend(x))].index
    data_day.iloc[weekend_rows, 3] = 1
    
    # 删除温度数据的第一行
    temperature.drop(index=0, axis=0, inplace=True)
    temperature.drop(['Date'], axis=1, inplace=True)
    temperature.reset_index(drop=True, inplace=True)
    
    # 添加温度列
    data_day = pd.concat([data_day, temperature], axis=1)
    return data_day

def prepare_data_1999():
    '''
    1999 年原始数据转换。
    '''
    # 1999 年数据
    file_path = r'../附件/原始数据/Load1999.xls'
    holiday_path = r'../附件/原始数据/Holidays.xls'
    temperature_path = r'../附件/原始数据/TemperatureJanuary1999.xls'
    # 读取数据
    data_day_1999 = pd.read_excel(file_path, index_col=0, parse_dates=True)
    # 按天进行下采样
    load_data_day = data_day_1999.iloc[:, :-1].astype('float').sum(axis=1)
    data_day = pd.DataFrame(columns = ['date'])
    
    data_day['date'] = load_data_day.index
    data_day['energy_load'] = load_data_day.values
    # 添加 是否假期 列
    holiday = pd.read_excel(holiday_path)
    # 删除第一列第一行（没用的数据）
    holiday.drop(index=0, axis=0, inplace=True)
    data_day['is_holiday'] = 0
    # 若数据的日期在假期表格内
    holiday_rows = data_day.loc[data_day['date'].isin(list(holiday.iloc[:,2]))].index
    data_day.iloc[holiday_rows, 2] = 1
    
    temperature = pd.read_excel(temperature_path)
    # 删除温度数据的第一行
    temperature.drop(index=0, axis=0, inplace=True)
    temperature.drop(['Day', 'Month'], axis=1, inplace=True)
    temperature.reset_index(drop=True, inplace=True)
    
    # 添加温度列
    data_day = pd.concat([data_day, temperature], axis=1)
    
    data_day['is_weekend'] =0
    weekend_rows = data_day.loc[data_day['date'].apply(lambda x: is_weekend(x))].index
    data_day.iloc[weekend_rows, 4] = 1
    
    # 保存 1999 年负荷数据
    data_day.to_excel(r'../附件/处理后的数据/天负荷数据_1999.xlsx')


def data_generator():
    '''
    由原始数据产生采样频率为天的数据
    '''
    # 1997 年数据
    file_path = r'../附件/原始数据/Load1997.xls'
    holiday_path = r'../附件/原始数据/Holidays.xls'
    temperature_path = r'../附件/原始数据/Temperature1997.xls'  
    # 读取数据
    data = pd.read_excel(file_path)
    holiday = pd.read_excel(holiday_path)
    temperature = pd.read_excel(temperature_path)
    # 得到 日期（天）-负荷-是否节假日-温度数据
    data_day_1997 = prepare_data(data, holiday, temperature)

    # 1998 年数据
    file_path = r'../附件/原始数据/Load1998.xls'
    holiday_path = r'../附件/原始数据/Holidays.xls'
    temperature_path = r'../附件/原始数据/Temperature1998.xls'
    # 读取数据
    data = pd.read_excel(file_path)
    temperature = pd.read_excel(temperature_path)
    holiday = pd.read_excel(holiday_path)
    # 得到 日期（天）-负荷-是否节假日-温度数据
    data_day_1998 = prepare_data(data, holiday, temperature, year=1998)    

    
    # 将 1997 年 和 1998 年数据存储起来
    # 两年的总数据
    data_day = pd.concat([data_day_1997, data_day_1998])
    data_day.to_excel(r'../附件/处理后的数据/天负荷数据.xlsx')
    data_day_1997.to_excel(r'../附件/处理后的数据/天负荷数据_1997.xlsx')
    data_day_1998.to_excel(r'../附件/处理后的数据/天负荷数据_1998.xlsx')
    
    prepare_data_1999()
    return data_day


if  __name__ == '__main__':
    # 由原始数据产生采样频率为天的电力负荷数据
    data_day = data_generator()
    
    # 使用 Dickey-Fuller 检验，判断数据是否稳定：
    result = adfuller(data_day['energy_load'])
    print('ADF Statistic: %f' % result[0])
    print('p-value: %f' % result[1])
    print('Critical Values:')
    for key, value in result[4].items():
        print('\t%s: %.3f' % (key, value))
    
    temperature = data_day['Temperature [oC]'].values
    data = data_day['energy_load'].values
    describe = data_day.describe()
    describe.to_excel(r'../附件/处理后的数据/天负荷数据_统计描述.xlsx')
    # 画出时间、电力负荷散点图、温度。
    plot_scatter(data, temperature)
    # 画概率密度图
    plot_hist(data)
    # 画自相关图
    plot_autocorrelation(data)
    
